//-------------------------------------------------------------
// Purpose: This program simulates a connect-four game.
//          The choice of moves is based on the an analysis
//          of potentially winning/losing sequences.
// Author:  John Gauch
//-------------------------------------------------------------
#include <cmath>
#include <cstdlib>
#include <iomanip>
#include <iostream>
#include <string>
using namespace std;

const bool DEBUG = true;
const char EMPTY = ' ';
const char BLACK = 'B';
const char WHITE = 'W';
const int LEN = 4;
const int ROWS = 6;
const int COLS = 7;
char this_player = EMPTY;
char board[ROWS][COLS];
int filled[COLS];
int move_count = 0;
int black_weight[LEN][LEN];
int white_weight[LEN][LEN];

//-------------------------------------------------------------
void init_weights(int weight[LEN][LEN], float base, float scale)
{
   // set all weights to -1
   for (int i = 0; i < LEN; i++)
      for (int j = 0; j < LEN; j++)
         weight[i][j] = -1;

   // set all black and all white weights
   // large base gives bias to long sequences
   // small base gives bias to short sequences
   // scale > 1 gives bias to offense
   // scale < 1 gives bias to defense
   for (int i = 0; i < LEN; i++)
   {
      weight[i][0] = pow(base,i+1);
      weight[0][i] = pow(base,i+1)*scale;
   }

   // print weights
   if (DEBUG)
   for (int i = 0; i < LEN; i++)
   {
      for (int j = 0; j < LEN; j++)
         cout << setw(6) << weight[i][j] << " ";
      cout << endl;
   }
}

//-------------------------------------------------------------
void init_board(char value)
{
   // loop over all rows and cols
   for (int c = 0; c < COLS; c++)
   for (int r = 0; r < ROWS; r++)
      board[r][c] = value;
   for (int c = 0; c < COLS; c++)
      filled[c] = 0;
}

//-------------------------------------------------------------
bool check_move(int col)
{
   // check location is valid
   return (col >= 0) && (col < COLS) && (filled[col] < ROWS);
}

//-------------------------------------------------------------
int score_sequence(int row, int col, int drow, int dcol, char player)
{
   // count black/white/total pieces
   int black_count = 0;
   int white_count = 0;
   int total_count = 0;
   for (int step = 0; step < LEN; step++)
   {
      int r = row + step * drow;
      int c = col + step * dcol;
      if ((r >= 0) && (r < ROWS) && (c >= 0) && (c < COLS))
      {
         if (board[r][c] == BLACK) black_count++;
         if (board[r][c] == WHITE) white_count++;
         total_count++;
      }
   }

   // calculate score
   if ((total_count == LEN) && (player == WHITE))
      return white_weight[black_count][white_count];
   else if ((total_count == LEN) && (player == BLACK))
      return black_weight[white_count][black_count];
   else
      return -1;
}

//-------------------------------------------------------------
int score_move(int col, char player)
{
   int score = random() % 10;
   if (check_move(col))
   {
      // check all offsets
      int row = filled[col];
      for (int step = 0; step < LEN; step++)
      {
         // check all directions
         score += score_sequence(row-step, col, 1, 0, player); 
         score += score_sequence(row, col-step, 0, 1, player); 
         score += score_sequence(row-step, col-step,  1, 1, player); 
         score += score_sequence(row+step, col-step, -1, 1, player); 
      }
   }
   return score; 
}

//-------------------------------------------------------------
void find_move(int &col, char player)
{
   // loop over all starting positions
   col = 0;
   int best_score = -1000;
   for (int c = 0; c < COLS; c++)
   if (check_move(c))
   {
      // calculate best score
      int score = score_move(c, player);
      if (DEBUG) cout << c << " " << score << endl;
      if (score > best_score)
      {
         col = c;
         best_score = score;
      }
   }
}

//-------------------------------------------------------------
void make_move(int col, char player)
{
   // make move
   if (check_move(col))
   {
      int row = filled[col]++;
      board[row][col] = player; 
      move_count++;
   }
}

//-------------------------------------------------------------
void showboard()
{
   // print all rows and cols
   cout << "\nMove: " << move_count << endl;
   for (int r = ROWS-1; r >= 0; r--)
   {
      cout << "+";
      for (int c = 0; c < COLS; c++)
         cout << "---+";
      cout << "\n";
      cout << "| ";
      for (int c = 0; c < COLS; c++)
         cout << board[r][c] << " | ";
      cout << "\n";
   }
   cout << "+";
   for (int c = 0; c < COLS; c++)
      cout << "---+";
   cout << "\n";
}

//-------------------------------------------------------------
bool check_win(char player)
{
   // loop over all starting positions
   for (int c = 0; c < COLS; c++)
   for (int r = 0; r < ROWS; r++)
   if (board[r][c] == player)
   {
      // check row
      int count = 0;
      for (int d = 0; d < LEN; d++)
         if ((r+d < ROWS) && (board[r+d][c] == player)) count++;
      if (count == LEN) return true;
    
      // check col
      count = 0;
      for (int d = 0; d < LEN; d++)
         if ((c+d < COLS) && (board[r][c+d] == player)) count++;
      if (count == LEN) return true;
    
      // check diagonal
      count = 0;
      for (int d = 0; d < LEN; d++)
         if ((r+d < ROWS) && (c+d < COLS) && (board[r+d][c+d] == player)) count++;
      if (count == LEN) return true;
    
      // check diagonal
      count = 0;
      for (int d = 0; d < LEN; d++)
         if ((r-d >= 0) && (c+d < COLS) && (board[r-d][c+d] == player)) count++;
      if (count == LEN) return true;
   }
   return false;
}

//-------------------------------------------------------------
void play_game()
{
   // initialize game
   int col = 0;
   move_count = 0;
   init_board(EMPTY);

   // make first move
   if (this_player == BLACK)
   {
      find_move(col, BLACK);
      make_move(col, BLACK);
      showboard();
      if (check_win(BLACK)) { cout << "BLACK wins\n"; return; }
   }

   // loop until someone wins game
   while (move_count < ROWS*COLS)
   {
      // white moves
      find_move(col, WHITE);
      make_move(col, WHITE);
      showboard();
      if (check_win(WHITE)) { cout << "WHITE wins\n"; return; }

      // black moves
      find_move(col, BLACK);
      make_move(col, BLACK);
      showboard();
      if (check_win(BLACK)) { cout << "BLACK wins\n"; return; }
   }

   // print final message
   cout << "Sorry, no one wins\n";
}

//-------------------------------------------------------------
int main(int argc, char *argv[])
{
   // process command arguments
   float b_base = 10;
   float b_scale = 2;
   float w_base = 10;
   float w_scale = 2;
   if (argc >= 2)
   {
      if ((argv[1][0] == 'b') || (argv[1][0] == 'B'))
         this_player = BLACK;
      if ((argv[1][0] == 'w') || (argv[1][0] == 'W'))
         this_player = WHITE;
   }
   if (argc == 6)
   {
      b_base = atof(argv[2]);
      b_scale = atof(argv[3]);
      w_base = atof(argv[4]);
      w_scale = atof(argv[5]);
   }
   
   // play game
   srandom(time(NULL));
   if (DEBUG) cout << "this_player " << this_player << endl;
   if (DEBUG) cout << "black weights\n";
   init_weights(black_weight, b_base, b_scale);
   if (DEBUG) cout << "white weights\n";
   init_weights(white_weight, w_base, w_scale);
   play_game();
   return 0;
}